package utils

import "crypto/sha256"

/**
author crastlin@163.com
date 2023-07-10
*/
import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"errors"
)

type AesEncryption struct {
	secret string
	iv     string
}

func NewAes(secret string, parameters ...string) (*AesEncryption, error) {
	aesStruct := &AesEncryption{}
	if secret == "" {
		return aesStruct, errors.New("请配置AES密钥")
	}
	aesStruct.secret = secret
	if len(parameters) == 0 || parameters[0] == "" {
		aesStruct.iv = aesStruct.secret[5:21]
	} else {
		aesStruct.iv = parameters[0]
	}
	return aesStruct, nil
}

func (a *AesEncryption) encryptSha256(key []byte) []byte {
	h := sha256.New()
	h.Write(key)
	return h.Sum(nil)
}

func (a *AesEncryption) PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padText...)
}

func (a *AesEncryption) PKCS7UnPadding(origData []byte) []byte {
	length := len(origData)
	unPadding := int(origData[length-1])
	return origData[:(length - unPadding)]
}

func (a *AesEncryption) Encrypt(plaintext []byte) ([]byte, error) {

	block, err := aes.NewCipher(a.encryptSha256([]byte(a.secret)))
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	plaintext = a.PKCS7Padding(plaintext, blockSize)
	ciphertext := make([]byte, len(plaintext))
	mode := cipher.NewCBCEncrypter(block, []byte(a.iv))
	mode.CryptBlocks(ciphertext, plaintext)

	return ciphertext, nil
}

func (a *AesEncryption) Decrypt(ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(a.encryptSha256([]byte(a.secret)))
	if err != nil {
		return nil, err
	}
	plaintext := make([]byte, len(ciphertext))
	mode := cipher.NewCBCDecrypter(block, []byte(a.iv))
	mode.CryptBlocks(plaintext, ciphertext)
	plaintext = a.PKCS7UnPadding(plaintext)

	return plaintext, nil
}
